import os
import time
import multiprocessing
from multiprocessing import Pool 


# 生成提问时的prompt
def run_rompt(input_table, output_json):
    cmd = (
        "python ./cluster_prompt.py -input " + input_table + 
        " -output " + output_json
    )
    os.system(cmd)
    

if __name__ == "__main__":
    start_time = time.time()
    # 生成任务列表
    files = os.listdir("03.cluster_add_pfam_info")
    file_path = [os.path.join("03.cluster_add_pfam_info",x) for x in files]
    output_path = [x.split("___")[0] + ".json"  for x in files]
    output_path = [os.path.join("./04.prompt",x) for x in output_path]
    os.makedirs("./04.prompt",exist_ok=True)
    task_list = list(zip(file_path, output_path))
    
    # 多进程执行任务
    cpu_count = multiprocessing.cpu_count() - 1
    with Pool(20) as pool:
        pool.starmap(run_rompt, task_list)
    end_time = time.time()
    total_time = end_time - start_time 
    print(f"\n任务数量是 {len(task_list)} :\n", f"\t用时 {total_time}s\n")
